use crate::common::packages::keepalive::KeepaliveReq;
use crate::common::packages::lsredir::{LsPortRedirReq, LsPortRedirResp};
use crate::common::packages::{FromSized, OS, PackageMetadata, PortRedirStatus, Role, SetGetRole};
use crate::common::quic::{QuicConn, QuicStream};
use anyhow::{Error, Result};
use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::sync::atomic::{AtomicI64, AtomicU8, Ordering};
use std::sync::{Arc, OnceLock};
use tokio::net::TcpListener;
use tokio::select;
use tokio::sync::RwLock;
use tokio::sync::broadcast::Sender as BroadcastSender;
use uuid::Uuid;

/// 记录uuid和被控端实例的映射关系
pub static CONTROLLED_INFO: OnceLock<RwLock<HashMap<Uuid, Arc<RwLock<ControlledInstance>>>>> =
    OnceLock::new();

/// 记录公网端口与被控端uuid的映射关系，key为公网监听端口，value为对应的被控端的uuid
pub static PUB_REDIR_MAP: OnceLock<Arc<RwLock<HashMap<u16, Uuid>>>> = OnceLock::new();

pub fn init() {
    CONTROLLED_INFO.get_or_init(|| RwLock::new(HashMap::new()));
    PUB_REDIR_MAP.get_or_init(|| Arc::new(RwLock::new(HashMap::new())));
}

/// 被控端实例结构体
pub struct ControlledInstance {
    uuid: Uuid,                                        // uuid
    os: Option<OS>,                                    // 操作系统
    private_ip: Option<Ipv4Addr>,                      // 私有IP地址
    addr: SocketAddr,                                  // 被控端的udp地址
    conn: Option<Arc<RwLock<QuicConn>>>,               // 连接被控端的quic连接
    keepalive_stream: Option<Arc<RwLock<QuicStream>>>, // 保活stream

    /// 公网端口转发信息,key 公网端口 val 端口转发实例
    port_map_info: Option<RwLock<HashMap<u16, Arc<RwLock<PortRedirInstance>>>>>,
}

impl ControlledInstance {
    pub fn from_cfg() -> Self {
        let cfg = super::GLOBAL_CFG.get().unwrap();
        let os = OS::os();
        ControlledInstance {
            conn: None,
            keepalive_stream: None,
            addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, cfg.quic_port)),
            uuid: cfg.uuid,
            os: Some(os),
            private_ip: None,
            port_map_info: None,
        }
    }

    pub fn is_proxy(&self) -> bool {
        self.uuid == super::GLOBAL_CFG.get().unwrap().uuid
    }

    pub async fn new(
        conn: Arc<RwLock<QuicConn>>,
        keepalive_stream: QuicStream,
    ) -> Arc<RwLock<Self>> {
        let peer_addr = conn.read().await.peer_addr();
        Arc::new(RwLock::new(ControlledInstance {
            conn: Some(conn),
            keepalive_stream: Some(Arc::new(RwLock::new(keepalive_stream))),
            addr: peer_addr,
            uuid: Uuid::nil(),
            port_map_info: None,
            os: None,
            private_ip: None,
        }))
    }

    /// 新建一个端口转发实例
    pub async fn new_redir_instance(
        &mut self,
        src_port: u16,
        dest_addr: SocketAddrV4,
    ) -> Result<Arc<RwLock<PortRedirInstance>>> {
        let port_info = super::get_port_config();
        if let Some(t) = port_info.as_ref() {
            let mut rw = t.write().await;
            if rw.can_add(src_port) {
                rw.add(src_port)?;
                drop(rw);
                let redir_instance = PortRedirInstance::new(self.uuid, dest_addr, src_port);
                if self.port_map_info.is_none() {
                    self.port_map_info = Some(RwLock::new(HashMap::new()));
                }
                let mut map_rw = self.port_map_info.as_ref().unwrap().write().await;
                map_rw.insert(src_port, redir_instance.clone());
                Ok(redir_instance)
            } else {
                Err(Error::msg("port has used"))
            }
        } else {
            Err(Error::msg("no port found"))
        }
    }

    /// 保活方法，会一直阻塞直到连接断开
    pub async fn keepalive(instance: Arc<RwLock<Self>>, init_md: PackageMetadata) -> Result<()> {
        let mut i = instance.write().await;
        let stream = i.keepalive_stream.clone();
        if stream.is_none() {
            return Err(Error::msg("keepalive stream is None"));
        }
        let stream = stream.unwrap();
        let mut pkg = KeepaliveReq::default();
        let mut stream_lock = stream.write().await;
        if let Err(e) = stream_lock.read_pkg_with_metadata(&init_md, &mut pkg).await {
            return Err(e);
        }
        if pkg.stat.len() != 1 {
            return Err(Error::msg("invalid keepalive stat"));
        }
        i.uuid = pkg.stat[0].uuid;
        i.os = Some(pkg.stat[0].os);
        i.private_ip = Some(pkg.stat[0].private_ip);
        drop(i);
        // 将实例存放到map中
        let mut m = CONTROLLED_INFO.get().unwrap().write().await;
        if m.contains_key(&pkg.stat[0].uuid) {
            m.remove(&pkg.stat[0].uuid).unwrap();
            m.insert(pkg.stat[0].uuid, instance.clone());
        } else {
            m.insert(pkg.stat[0].uuid, instance.clone());
        }
        drop(m);
        let uuid = pkg.stat[0].uuid;
        let mut resp = pkg.to_resp(Role::Proxy);
        drop(pkg);
        let mut new_pkg = KeepaliveReq::default();

        async fn remove_instance(uuid: &Uuid) {
            CONTROLLED_INFO.get().unwrap().write().await.remove(uuid);
            // todo
        }

        let err: Error;

        // 循环，应答保活流
        loop {
            if let Err(e) = stream_lock.read_pkg(&mut new_pkg).await {
                err = Error::from(e);
                break;
            }
            if new_pkg.stat.len() == 1 {
                err = Error::msg("invalid keepalive stat");
                break;
            }
            if new_pkg.get_role() != Role::Controlled || new_pkg.stat[0].uuid != uuid {
                err = Error::msg("invalid keepalive stat");
                break;
            }
            resp.cmd_uuid = new_pkg.cmd_uuid;
            if let Err(e) = stream_lock.write_pkg(&resp).await {
                err = Error::from(e);
                break;
            }
        }
        remove_instance(&uuid).await;
        Err(err)
    }

    /// 关闭被控端实例
    pub fn down(instance: Arc<RwLock<Self>>) -> Result<()> {
        /*
            关闭相关的转发实例
            清除PUB_REDIR_MAP中的相关信息
            清除PORT_CONFIG中的相关信息
        */
        Ok(())
    }

    pub async fn get_port_map_info(&self) -> Option<HashMap<u16, ()>> {
        match self.port_map_info.as_ref() {
            Some(map) => {
                todo!()
            },
            None => None,
        }
    }
}

/// 公网端口转发实例
pub struct PortRedirInstance {
    controlled_uuid: Uuid,             // 被控端的uuid
    dest_addr: SocketAddrV4,           // 端口转发的目的地址
    src_port: u16,                     // 端口转发的源端口
    status: AtomicU8,                  // 端口转发状态
    tcp_listener: Option<TcpListener>, // tcp监听器
    connect_num: Arc<AtomicI64>,       // 该公网转发实例的连接数
    broadcast: BroadcastSender<()>,    // 强制关闭通知器
}

impl PortRedirInstance {
    pub fn new(uuid: Uuid, dest_addr: SocketAddrV4, src_port: u16) -> Arc<RwLock<Self>> {
        let (s, _) = tokio::sync::broadcast::channel::<()>(10);
        Arc::new(RwLock::new(PortRedirInstance {
            controlled_uuid: uuid,
            dest_addr,
            src_port,
            status: AtomicU8::new(PortRedirStatus::HardStop(true).into()),
            tcp_listener: None,
            connect_num: Arc::new(AtomicI64::new(0)),
            broadcast: s,
        }))
    }

    /// 开始监听
    pub async fn listen(&mut self) -> Result<()> {
        // 只有当现在的状态是完全关闭时，才能打开
        if let Ok(_) = self.status.compare_exchange(
            PortRedirStatus::HardStop(true).into(),
            PortRedirStatus::Open(true).into(),
            Ordering::SeqCst,
            Ordering::SeqCst,
        ) {
            let listener = TcpListener::bind(("0.0.0.0", self.src_port)).await?;
            self.tcp_listener = Some(listener);
            let open: u8 = PortRedirStatus::Open(true).into();
            while self.status.load(Ordering::SeqCst) == open {
                let (mut stream, _) = self.tcp_listener.as_ref().unwrap().accept().await?;
                let conn_num = self.connect_num.clone();
                if let Some(conn) = CONTROLLED_INFO
                    .get()
                    .unwrap()
                    .write()
                    .await
                    .get(&self.controlled_uuid)
                {
                    let c = conn.read().await.conn.as_ref().unwrap().clone();
                    conn_num.fetch_sub(1, Ordering::SeqCst);
                    let mut r = self.broadcast.subscribe();
                    tokio::spawn(async move {
                        let c = QuicConn::open_bi(c).await;
                        let mut s = c.unwrap();
                        select! {
                            _ = s.redir(&mut stream) => {


                            }
                            _ = r.recv() => {
                                // 接收到强制关闭的信号
                            }
                        }
                        s.close_stream(0).await;
                        conn_num.fetch_add(-1, Ordering::SeqCst);
                    });
                };
            }
            Ok(())
        } else {
            Err(Error::msg("port is exhausted"))
        }
    }

    pub async fn close(&mut self, state: PortRedirStatus) -> Result<()> {
        let stat_now = match PortRedirStatus::from_sized(self.status.load(Ordering::SeqCst)) {
            Ok((s, _)) => s,
            Err(e) => {
                return Err(e);
            }
        };
        match stat_now {
            PortRedirStatus::Open(_) => match state {
                PortRedirStatus::Open(_) => Err(Error::msg("invalid target state")),
                PortRedirStatus::SoftStop(_) => {
                    self.status.store(state.into(), Ordering::SeqCst);
                    Ok(())
                }
                PortRedirStatus::HardStop(_) => {
                    self.status.store(state.into(), Ordering::SeqCst);
                    if state == PortRedirStatus::HardStop(true) {
                        self.broadcast.send(())?;
                    }
                    Ok(())
                }
            },
            PortRedirStatus::SoftStop(_) => match state {
                PortRedirStatus::Open(_) | PortRedirStatus::SoftStop(_) => {
                    Err(Error::msg("invalid target state"))
                }
                PortRedirStatus::HardStop(_) => {
                    self.status.store(state.into(), Ordering::SeqCst);
                    if state == PortRedirStatus::HardStop(true) {
                        self.broadcast.send(())?;
                    }
                    Ok(())
                }
            },
            PortRedirStatus::HardStop(_) => Ok(()),
        }
    }
}

/*
    被控端处理的请求：
    CloseConnReq
    PortRedirReq
    ShellCmdReq
    Socks5RedirReq

    缓存需要处理的请求：
    ControlledReq
    LsPortRedirReq
*/

// todo CloseConnReq

// todo PortRedirReq

// todo ShellCmdReq

// todo Socks5RedirReq

// todo LsPortRedirReq

pub async fn list_port_redir(req: &LsPortRedirReq) -> Result<LsPortRedirResp> {
    let rw_port = super::PORT_CONFIG.get().unwrap().as_ref();
    match rw_port {
        None => Ok(LsPortRedirResp::new(Role::Proxy, req.cmd_uuid, false, 0, 0)),
        Some(p) => {
            let rw_port = p.as_ref().read().await;
            let r = LsPortRedirResp::new(Role::Proxy, req.cmd_uuid, true, rw_port.from, rw_port.to);
            drop(rw_port);
            let maps = PUB_REDIR_MAP.get().unwrap().read().await.keys();
            let info = CONTROLLED_INFO.get().unwrap().read().await.keys();

            Ok(r)
        }
    }
}

// todo ControlledReq
